# Copyright 2025 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Templates for code auto generation."""
import re


class Template:
    """
    template for generate c++/python code
    """
    regular_str = r"(^[^\n\S]*)?\$([^\d\W]\w*|\{,?[^\d\W]\w*\,?})"
    regular_match = re.compile(regular_str, re.MULTILINE)

    def __init__(self, code_pattern):
        self.code_pattern = code_pattern

    @staticmethod
    def load_from_file(file_path):
        """load template from file"""
        with open(file_path, "r") as f:
            return Template(f.read())

    def replace(self, **kwargs):
        """
        replace param.
        :param kwargs:
        :return:
        """

        def find(key: str):
            if key in kwargs:
                return kwargs[key]
            raise TypeError(f"{key} should be in kwargs!")

        def add_indent(indent, var):
            return "".join([indent + line + "\n" for data in var for line in str(data).splitlines()]).rstrip()

        def extract_variable(key):
            start = ""
            end = ""
            if key[0] == "{":
                key = key[1:-1]
                if key[0] == ",":
                    start = ","
                    key = key[1:]
                if key[-1] == ",":
                    end = ", "
                    key = key[:-1]
            return find(key), start, end

        def match_rule(match):
            indent = match.group(1)
            key = match.group(2)
            var, start, end = extract_variable(key)
            if indent is not None:
                if not isinstance(var, list):
                    return add_indent(indent, [var])
                return add_indent(indent, var)
            if isinstance(var, list):
                code = ", ".join(str(x) for x in var)
                if not var:
                    return code
                return start + code + end
            return str(var)

        return self.regular_match.sub(match_rule, self.code_pattern)


NEW_LINE = "\n"

PYTHON_PRIM_TEMPLATE = Template("""

class _Pyboost${class_name}Prim(${class_name}Prim_):
    def __call__(self, ${input_args}):
        ${process_func}
        return super().__call__([${processed_args}])


${func_impl_name}_impl = _Pyboost${class_name}Prim()
""")

IMPORT_PYBOOST_PRIM_HEADER = f"""
from mindspore.ops._utils.arg_handler import *
"""

IMPORT_PYBOOST_FUNC_HEADER = f"""
from mindspore.common import dtype as mstype
from mindspore.ops.auto_generate.pyboost_inner_prim import *

"""

PY_LICENSE_STR = f"""# Copyright 2023 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""

OPS_PY_PRIM_HEADER = f"""
\"\"\"Operators definition generated by gen_ops.py, includes primitive classes.\"\"\"

from mindspore.ops.primitive import Primitive, prim_arg_register
from mindspore.ops import signature as sig
from mindspore.common import dtype as mstype
from mindspore.common._decorator import deprecated
from mindspore.ops._primitive_cache import _get_cache_prim
from mindspore.ops._utils.arg_dtype_cast import type_it
from mindspore.ops._utils.arg_handler import *
from mindspore._c_expression import OpDtype
from mindspore.common.jit_context import jit_context
from mindspore._checkparam import is_stub_tensor
"""

OPS_PY_DEF_HEADER = f"""
\"\"\"Operators definition generated by gen_ops.py, includes functions.\"\"\"

from .gen_ops_prim import *
from .pyboost_inner_prim import *
from mindspore.ops.operations.manually_defined.ops_def import *
from mindspore.ops._primitive_cache import _get_cache_prim
"""

CUSTOM_OPS_PY_DEF_HEADER = f"""
\"\"\"Operators definition generated by gen_ops.py, includes functions.\"\"\"

from .gen_ops_prim import *
"""

PRIMITIVE_CLASS_DESC = """    r\"\"\"
    .. code-block::

        prim = ops.$class_name($init_args_str)
        out = prim($input_args_str)

    is equivalent to

    .. code-block::

        ops.$func_name($args_str)

    Refer to :func:`mindspore.ops.$func_name` for more details.
    \"\"\"
"""

CC_LICENSE_STR = f"""/**
 * Copyright 2023-2025 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */"""

OP_PROTO_TEMPLATE = Template("""
${func_impl_declaration}
OpDef g${class_name} = {
  /*.name_=*/"${class_name}",
  /*.args_=*/ {
    ${input_args}
  },
  /* .returns_ = */ {
    ${return_args}
  },
  /*.signatures_ =*/ {
    ${signatures}
  },
  /*.indexes_ =*/ {
    ${indexes}
  },
  /*.func_impl_=*/${func_impl_define},
  /*.enable_dispatch_ =*/${enable_dispatch},
  /*.is_view_ =*/${is_view},
  /*.is_graph_view_ =*/${is_graph_view},
};
REGISTER_PRIMITIVE_OP_DEF(${class_name}, &g${class_name});
""")

OP_PRIM_CLASS_DEFINE_TEMPLATE = Template("""

class ${class_name}(Primitive):
${class_desc}${signature_code}${deprecated_code}
${init_method}

${call_method}""")
